"""
MySQL与PostgreSQL数据一致性校验工具
功能：全量对比表结构、数据行差异，输出详细报告
"""

import hashlib
import logging
from dataclasses import dataclass
from typing import Dict, List, Tuple
import mysql.connector
import psycopg2
from mysql.connector import errorcode
from psycopg2 import OperationalError

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.FileHandler('db_compare.log'), logging.StreamHandler()]
)


@dataclass
class DBConfig:
    """数据库连接配置"""
    mysql_host: str
    mysql_port: int
    mysql_user: str
    mysql_password: str
    mysql_database: str
    pg_host: str
    pg_port: int
    pg_user: str
    pg_password: str
    pg_database: str


class DatabaseComparator:
    def __init__(self, config: DBConfig):
        self.config = config
        self.mysql_conn = None
        self.pg_conn = None

    def _connect_mysql(self) -> bool:
        """连接MySQL数据库"""
        try:
            self.mysql_conn = mysql.connector.connect(
                host=self.config.mysql_host,
                port=self.config.mysql_port,
                user=self.config.mysql_user,
                password=self.config.mysql_password,
                database=self.config.mysql_database,
                charset='utf8mb4',
                connection_timeout=10
            )
            return True
        except mysql.connector.Error as err:
            logging.error(f"MySQL连接失败: {err}")
            return False

    def _connect_pg(self) -> bool:
        """连接PostgreSQL数据库"""
        try:
            self.pg_conn = psycopg2.connect(
                host=self.config.pg_host,
                port=self.config.pg_port,
                user=self.config.pg_user,
                password=self.config.pg_password,
                dbname=self.config.pg_database,
                connect_timeout=10
            )
            return True
        except OperationalError as err:
            logging.error(f"PostgreSQL连接失败: {err}")
            return False

    def _get_table_schemas(self) -> Tuple[Dict, Dict]:
        """获取表结构信息"""
        # 获取MySQL表结构
        mysql_schemas = {}
        mysql_cursor = self.mysql_conn.cursor(dictionary=True)
        mysql_cursor.execute("""
            SELECT TABLE_NAME, COLUMN_NAME, DATA_TYPE, COLUMN_TYPE 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_SCHEMA = %s
        """, (self.config.mysql_database,))
        for row in mysql_cursor:
            table = row['TABLE_NAME']
            if table not in mysql_schemas:
                mysql_schemas[table] = []
            mysql_schemas[table].append({
                'column': row['COLUMN_NAME'],
                'type': row['DATA_TYPE'],
                'full_type': row['COLUMN_TYPE']
            })

        # 获取PostgreSQL表结构
        pg_schemas = {}
        pg_cursor = self.pg_conn.cursor()
        pg_cursor.execute("""
            SELECT table_name, column_name, data_type, udt_name
            FROM information_schema.columns
            WHERE table_schema = 'public'
        """)
        for row in pg_cursor:
            table = row
            if table not in pg_schemas:
                pg_schemas[table] = []
            pg_schemas[table].append({
                'column': row,
                'type': row,
                'udt': row  # 处理自定义类型如int4, timestamp等
            })

        return mysql_schemas, pg_schemas

    def _type_mapping(self, mysql_type: str, pg_udt: str) -> bool:
        """类型兼容性检查"""
        type_map = {
            'varchar': 'character varying',
            'text': 'text',
            'int': 'integer',
            'bigint': 'bigint',
            'datetime': 'timestamp',
            'tinyint': 'boolean' if mysql_type == 'tinyint(1)' else 'smallint'
        }
        expected_pg_type = type_map.get(mysql_type.split('('), 'text')
        return expected_pg_type == pg_udt

    def compare_schemas(self) -> Dict:
        """结构比对核心逻辑"""
        mysql_schemas, pg_schemas = self._get_table_schemas()
        differences = {}

        # 检查缺失表
        for table in mysql_schemas:
            if table not in pg_schemas:
                differences[table] = {'status': 'missing_in_pg'}

        # 检查字段差异
        for table, columns in pg_schemas.items():
            if table not in mysql_schemas:
                differences[table] = {'status': 'missing_in_mysql'}
                continue

            mysql_cols = {str.lower(col['column']): col for col in mysql_schemas[table]}
            for pg_col in columns:
                # 这里需要忽略大小写，所以全部转小写后再对比
                col_name = str.lower(pg_col['column'])
                if col_name not in mysql_cols:
                    differences.setdefault(table, {}).setdefault('extra_columns', []).append(col_name)
                    continue

                # 类型兼容检查
                mysql_type = mysql_cols[col_name]['full_type']
                if not self._type_mapping(mysql_type, pg_col['udt']):
                    differences.setdefault(table, {}).setdefault('type_mismatch', []).append({
                        'column': col_name,
                        'mysql': mysql_type,
                        'pg': pg_col['udt']
                    })

        return differences

    def _generate_row_hash(self, row: Tuple, col_types: List) -> str:
        """生成行哈希（处理特殊类型）"""
        hash_str = ''
        for value, col_type in zip(row, col_types):
            # 处理二进制数据
            if isinstance(value, (bytes, bytearray)):
                hash_str += hashlib.sha256(value).hexdigest()
            # 处理日期时间
            elif 'datetime' in col_type or 'timestamp' in col_type:
                hash_str += str(value.timestamp()) if value else 'NULL'
            # 处理数值型
            elif 'int' in col_type or 'float' in col_type:
                hash_str += f"{value:.5f}" if isinstance(value, float) else str(value)
            else:
                hash_str += str(value) if value is not None else 'NULL'
        return hashlib.sha256(hash_str.encode()).hexdigest()

    def compare_table_data(self, table: str, pk_column: str = 'id', batch_size: int = 1000) -> Dict:
        """数据比对核心逻辑"""
        # 获取主键范围
        mysql_cursor = self.mysql_conn.cursor()
        mysql_cursor.execute(f"SELECT MIN({pk_column}), MAX({pk_column}) FROM {table}")
        min_id, max_id = mysql_cursor.fetchone()

        pg_cursor = self.pg_conn.cursor()
        pg_cursor.execute(f"SELECT MIN({pk_column}), MAX({pk_column}) FROM {table}")
        pg_min, pg_max = pg_cursor.fetchone()

        if (min_id, max_id) != (pg_min, pg_max):
            return {'error': '主键范围不一致', 'mysql_range': (min_id, max_id), 'pg_range': (pg_min, pg_max)}

        # 获取字段类型
        mysql_cursor.execute(f"SHOW COLUMNS FROM {table}")
        col_types = [col for col in mysql_cursor.fetchall()]

        differences = []
        current_id = min_id

        # 分块比对
        while current_id <= max_id:
            end_id = current_id + batch_size
            # 从MySQL读取数据
            mysql_cursor.execute(
                f"SELECT * FROM {table} WHERE {pk_column} BETWEEN %s AND %s ORDER BY {pk_column}",
                (current_id, end_id)
            )
            mysql_rows = {row: row for row in mysql_cursor.fetchall()}

            # 从PG读取数据
            pg_cursor.execute(
                f"SELECT * FROM {table} WHERE {pk_column} BETWEEN %s AND %s ORDER BY {pk_column}",
                (current_id, end_id)
            )
            pg_rows = {row: row for row in pg_cursor.fetchall()}

            # 比较哈希
            for pk in range(current_id, end_id + 1):
                mysql_row = mysql_rows.get(pk)
                pg_row = pg_rows.get(pk)

                if not mysql_row and not pg_row:
                    continue
                if not mysql_row:
                    differences.append({'pk': pk, 'status': 'missing_in_mysql'})
                    continue
                if not pg_row:
                    differences.append({'pk': pk, 'status': 'missing_in_pg'})
                    continue

                mysql_hash = self._generate_row_hash(mysql_row, col_types)
                pg_hash = self._generate_row_hash(pg_row, col_types)

                if mysql_hash != pg_hash:
                    differences.append({
                        'pk': pk,
                        'mysql_hash': mysql_hash,
                        'pg_hash': pg_hash
                    })

            current_id = end_id + 1

        return {'total': len(differences), 'differences': differences}
